import numpy as np
import cv2
import gradio as gr
from sklearn.cluster import KMeans
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches


def extract_color_palette(image, k=5):
    """
    从图像中提取主色调调色板
    :param image: PIL Image 或 numpy array
    :param k: 提取的颜色数量
    :return: 调色板图像和颜色十六进制列表
    """
    # 转换为numpy array（如果是PIL Image）
    if isinstance(image, Image.Image):
        img_array = np.array(image)
    else:
        img_array = image
    
    # 确保是RGB格式
    if len(img_array.shape) == 2:  # 灰度图
        img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
    elif img_array.shape[2] == 4:  # RGBA
        img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2RGB)
    
    # 将图像重塑为二维矩阵（像素点 x RGB）
    pixels = img_array.reshape((-1, 3))
    
    # 使用KMeans进行聚类
    kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
    kmeans.fit(pixels)
    
    # 获取聚类中心（主色调）
    colors = kmeans.cluster_centers_.astype(int)
    
    # 计算每个颜色的像素数量（用于排序）
    labels = kmeans.labels_
    label_counts = np.bincount(labels)
    
    # 按像素数量排序颜色（从多到少）
    sorted_indices = np.argsort(-label_counts)
    colors = colors[sorted_indices]
    
    # 转换为十六进制
    hex_colors = ['#{:02x}{:02x}{:02x}'.format(r, g, b) for r, g, b in colors]
    
    # 创建调色板可视化
    palette_image = create_palette_visualization(colors, hex_colors)
    
    return palette_image, hex_colors


def create_palette_visualization(colors, hex_colors):
    """
    创建调色板可视化图像
    :param colors: RGB颜色数组
    :param hex_colors: 十六进制颜色列表
    :return: matplotlib图像
    """
    fig, ax = plt.subplots(figsize=(12, 3))
    
    # 绘制色块
    num_colors = len(colors)
    for i, (color, hex_color) in enumerate(zip(colors, hex_colors)):
        # 绘制矩形色块
        rect = patches.Rectangle((i, 0), 1, 1, 
                                 facecolor=np.array(color)/255.0,
                                 edgecolor='white',
                                 linewidth=2)
        ax.add_patch(rect)
        
        # 添加十六进制文本
        # 计算文本颜色（根据背景亮度选择黑色或白色）
        brightness = (color[0] * 299 + color[1] * 587 + color[2] * 114) / 1000
        text_color = 'white' if brightness < 128 else 'black'
        
        ax.text(i + 0.5, 0.5, hex_color.upper(), 
               ha='center', va='center',
               fontsize=14, fontweight='bold',
               color=text_color,
               family='monospace')
    
    # 设置坐标轴
    ax.set_xlim(0, num_colors)
    ax.set_ylim(0, 1)
    ax.axis('off')
    
    plt.tight_layout()
    
    # 转换为numpy array（兼容新旧版本的matplotlib）
    fig.canvas.draw()
    try:
        # 新版本matplotlib
        buf = fig.canvas.buffer_rgba()
        palette_image = np.asarray(buf)
        # 转换RGBA到RGB
        palette_image = palette_image[:, :, :3]
    except AttributeError:
        # 旧版本matplotlib
        palette_image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        palette_image = palette_image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    
    plt.close(fig)
    
    return palette_image


def process_image(image):
    """
    处理上传的图像并返回调色板
    :param image: 上传的图像
    :return: 调色板图像和颜色信息文本
    """
    if image is None:
        return None, "请上传一张图片"
    
    # 提取调色板
    palette_image, hex_colors = extract_color_palette(image, k=5)
    
    # 创建颜色信息文本
    color_info = "色彩十六进制表示形式：\n\n"
    for i, hex_color in enumerate(hex_colors, 1):
        color_info += f"{i}. {hex_color.upper()}\n"
    
    return palette_image, color_info


# 创建Gradio界面
with gr.Blocks(title="KMeans调色板提取器", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎨 KMeans调色板提取器
        ### 上传图像获取其主色调调色板（5个颜色）
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            # 输入区域
            image_input = gr.Image(
                label="上传图像获取调色板",
                type="pil",
                height=400
            )
            
            with gr.Row():
                clear_btn = gr.Button("Clear", variant="secondary")
                submit_btn = gr.Button("Submit", variant="primary")
        
        with gr.Column(scale=1):
            # 输出区域
            palette_output = gr.Image(
                label="调色板",
                height=150
            )
            
            color_text = gr.Textbox(
                label="色彩十六进制表示形式",
                lines=7,
                interactive=False
            )
            
            flag_btn = gr.Button("Flag", variant="secondary")
    
    # 事件处理
    submit_btn.click(
        fn=process_image,
        inputs=[image_input],
        outputs=[palette_output, color_text]
    )
    
    clear_btn.click(
        fn=lambda: (None, None, ""),
        inputs=[],
        outputs=[image_input, palette_output, color_text]
    )
    
    # 自动处理（上传后立即处理）
    image_input.change(
        fn=process_image,
        inputs=[image_input],
        outputs=[palette_output, color_text]
    )


# 启动应用
if __name__ == "__main__":
    demo.launch(share=False, server_name="127.0.0.1", server_port=7860)
